# 导入所需库
from transformers import AutoTokenizer, AutoModel
import faiss
import numpy as np
import torch

# 初始化模型和分词器（使用MiniLM模型）
model_name = "sentence-transformers/all-MiniLM-L4-v2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 文档数据集
documents = [
    "人工智能是一种模仿人类智能的技术。",
    "深度学习是一种基于神经网络的机器学习方法。",
    "自然语言处理使得计算机能够理解和生成文本。",
    "机器学习让计算机通过数据自动进行学习。",
    "数据科学帮助分析复杂数据以提取信息。"
]


# 将文本转化为向量
def embed_texts(texts):
    inputs = tokenizer(texts, padding=True, truncation=True,
                       return_tensors="pt")
    with torch.no_grad():
        embeddings = model(**inputs). \
            last_hidden_state.mean(dim=1)  # 平均池化获取句子嵌入
    return embeddings.cpu().numpy()


# 获取文档向量嵌入
doc_embeddings = embed_texts(documents)
# 创建FAISS索引
dimension = doc_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(doc_embeddings)


# 查询函数
def search(query, top_k=2):
    query_embedding = embed_texts([query])
    distances, indices = index.search(query_embedding, top_k)
    return [(documents[i], distances[0][j]) \
            for j, i in enumerate(indices[0])]


# 测试查询
query_text = "计算机如何理解文本？"
results = search(query_text)
print("查询结果：")
for result, distance in results:
    print(f"文档: {result}, 相似度距离: {distance:.4f}")
